#!/data/Separated_Users/zhangjuan/miniconda3/envs/R4/bin/Rscript
library(getopt)
spec <- matrix(
  c(
    "help",      "h",  0,  "logical",
    "rds",       "r",  1,  "character",
    "deg",       "d",  2,  "character",
    "ncpu",      "n",  2,  "numeric",
    "min_exp",   "M",  2,  "numeric",
    "mean_exp",  "E",  2,  "numeric",
    "pval",      "p",  2,  "numeric",
    "pval_adj",  "q",  2,  "numeric",
    "pseu_pval", "P",  2,  "numeric",
    "pseu_qval", "Q",  2,  "numeric",
    "var_gene",  "v",  1,  "numeric",
    "xgene",     "x",  2,  "character",
    "od",        "o",  1,  "character"
  ),
  byrow = TRUE, ncol = 4
)
opt <- getopt(spec)

print_usage <- function(spec = NULL) {
  cat("Usage example:

  Rscript monocle2.R  --rds  mydata.rds --od  ./ --var_gene 3

  Options:
    --help  -h 	 NULL 	      get this help
    --rds 	     character    rds file [forced]
    --deg	       character	  use clusters diff genes from Seurat
    --ncpu	     numeric	    the number of cpu to be used, default: 8
    --min_exp	   numeric	    the number of genes above this threshold are detectable in each cell
    --mean_exp	 numeric	    threshold of mean_exp used to monocle selected var genes
    --pval	     numeric	    threshold of pvalue used to select clusters diff genes
    --pval_adj   numeric	    threshold of pval_adj used to select clusters diff genes, default pval_adj: 0.001
    --pseu_pval  numeric	    threshold of pval used to select diff pseudotime-dependent genes
    --pseu_qval  numeric	    threshold of qval used to select diff pseudotime-dependent genes, default pseu_qval: 0.01
    --xgene      character    plot xgene heatmap

    --var_gene   numeric      a vector of feature gene used for ordering cells, can be: 1, 2, 3, default: 3
                                 1: use clusters diff genes of Seurat's result [if select 1, must privided --deg parameter]
                                 2: use seurat selected var genes
                                 3: use monocle selected var genes

    --od         character   outdir [forced]
  \n")
  q(status = 1)
}

# opt <- list()
# opt$rds <- "/Pub/Users/cuiye/projects/F240522001//results/3.1.subtype_res/object_x.rds"
# opt$var_gene <- 1
# opt$od <- "/Pub/Users/cuiye/projects/F240522001//results/4.1.subtype_monocle/"
# opt$deg <- "/Pub/Users/cuiye/projects/F240522001//results/2.1.epi_subtype/cluster_markers.txt"
# opt$xgene <- "/Pub/Users/cuiye/projects/F240522001//data/xgene.txt"
# opt$pval <- 0.05

if (!is.null(opt$help) | is.null(opt$rds) | is.null(opt$od)) {
  print_usage(spec)
}
if (is.null(opt$min_exp)) {
  opt$min_exp <- 0.1
}
if (is.null(opt$mean_exp)) {
  opt$mean_exp <- 0.1
}
if (is.null(opt$pval) & is.null(opt$pval_adj)) {
  opt$pval_adj <- 0.01
}
if (is.null(opt$pseu_pval) & is.null(opt$pseu_qval)) {
  opt$pseu_qval <- 0.01
}
if (is.null(opt$var_gene)) {
  opt$var_gene <- 3
}
if (is.null(opt$ncpu)) {
  opt$ncpu <- 8
}
if (!file.exists(opt$od)) {
  dir.create(opt$od, recursive = T)
}

source("/Pub/Users/cuiye/RCodes/UserCode/newlover/color_fun.R")
library(monocle)
library(Seurat)
library(ggplot2)
library(patchwork)
suppressMessages(library(tidyverse))

## read rds file
scRNA <- readRDS(opt$rds)
scRNA$celltype <- scRNA@active.ident

## 1.构造monocle需要的对象
data <- scRNA@assays$RNA$counts
data <- data[, match(rownames(scRNA@meta.data), colnames(data))]

fData <- data.frame(gene_short_name = row.names(data), row.names = row.names(data))
fd <- new("AnnotatedDataFrame", data = fData)
pd <- new("AnnotatedDataFrame", data = scRNA@meta.data)

mycds <- newCellDataSet(data, phenoData = pd, featureData = fd, expressionFamily = negbinomial.size())
mycds

## 2.标准化处理
mycds <- estimateSizeFactors(mycds)
mycds <- estimateDispersions(mycds, cores = opt$ncpu, relative_expr = TRUE)
mycds <- detectGenes(mycds, min_expr = opt$min_exp)

## 3.选择代表性基因
if (opt$var_gene == 1) { # 1.使用clusters差异表达基因

  diff.genes <- read.table(opt$deg, header = T, sep = "\t", comment.char = "", check.names = F)

  if (!is.null(opt$pval)) {
    ordering_genes <- subset(diff.genes, p_val < opt$pval)$gene
  } else if (!is.null(opt$pval_adj)) {
    ordering_genes <- subset(diff.genes, p_val_adj < opt$pval_adj)$gene
  } else {
    print("pval and p_val_adj must privided one!")
  }
}

if (opt$var_gene == 2) { # 2.使用seurat选择的高变基因
  ordering_genes <- VariableFeatures(scRNA)
}

if (opt$var_gene == 3) { # 3.使用monocle选择的高变基因
  disp_table <- dispersionTable(mycds)
  # ordering_genes <- subset(disp_table, mean_expression >= opt$mean_exp & dispersion_empirical >= 1 * dispersion_fit)$gene_id
  ordering_genes <- subset(disp_table, mean_expression >= opt$mean_exp)$gene_id
}

head(ordering_genes)
length(ordering_genes)

## 4.降维及细胞排序：使用ordering_genes开展后续分析
mycds <- setOrderingFilter(mycds, ordering_genes)
mycds <- reduceDimension(mycds, max_components = 2, method = "DDRTree")
mycds <- orderCells(mycds)

p1 <- plot_ordering_genes(mycds)
ggsave(filename = paste0(opt$od, "/plot_ordering_genes.pdf"), plot = p1, width = 5, height = 5)

# State轨迹分布图
plot1 <- plot_cell_trajectory(mycds, color_by = "State", cell_size = 1) + scale_color_manual(values = color_fun3)
ggsave(filename = paste0(opt$od, "/Pseudotime_State.pdf"), plot = plot1, width = 5, height = 5)

# Cluster轨迹分布图
plot2 <- plot_cell_trajectory(mycds, color_by = "celltype", cell_size = 1) + scale_color_manual(values = c(color_fun7, color_fun3, color_fun4))
ggsave(filename = paste0(opt$od, "/Pseudotime_celltype.pdf"), plot = plot2, width = 5, height = 5)

# Pseudotime轨迹图
plot3 <- plot_cell_trajectory(mycds, color_by = "Pseudotime", cell_size = 1)
ggsave(filename = paste0(opt$od, "/Pseudotime_time.pdf"), plot = plot3, width = 5, height = 5)

# 合并作图
plotc <- plot1 | plot2 | plot3
ggsave(filename = paste0(opt$od, "/Pseudotime_Combine.pdf"), plot = plotc, width = 15, height = 5)

# 保存结果
write.csv(pData(mycds), paste0(opt$od, "/Pseudotime.csv"))
saveRDS(mycds, paste0(opt$od, "/Pseudotime_mycds.rds"))

## 轨迹图分面显示
p1 <- plot_cell_trajectory(mycds, color_by = "State") + facet_wrap(~State, nrow = 1) + scale_color_manual(values = color_fun3)
p2 <- plot_cell_trajectory(mycds, color_by = "celltype") + facet_wrap(~celltype, nrow = 1) + scale_color_manual(values = c(color_fun7, color_fun3, color_fun4))
plotc <- p1 / p2
ggsave(filename = paste0(opt$od, "/Pseudotime_trajectory_facet.pdf"), plot = plotc, width = 15, height = 10)


## 5.Monocle基因可视化
myth <- theme(axis.text.x = element_text(angle = 45))
s.genes <- as.character(head(ordering_genes))
p1 <- plot_genes_jitter(mycds[s.genes, ], grouping = "State", color_by = "State") + myth + scale_color_manual(values = color_fun3)
p2 <- plot_genes_violin(mycds[s.genes, ], grouping = "State", color_by = "State") + myth + scale_fill_manual(values = color_fun3)
p3 <- plot_genes_in_pseudotime(mycds[s.genes, ], color_by = "State") + myth + scale_color_manual(values = color_fun3)
plotc <- p1 | p2 | p3
ggsave(filename = paste0(opt$od, "/Dispgenes_top6_visual_by_State.pdf"), plot = plotc, width = 8, height = 8)
##
p1 <- plot_genes_jitter(mycds[s.genes, ], grouping = "celltype", color_by = "celltype") + myth + scale_color_manual(values = c(color_fun7, color_fun3, color_fun4))
p2 <- plot_genes_violin(mycds[s.genes, ], grouping = "celltype", color_by = "celltype") + myth + scale_fill_manual(values = c(color_fun7, color_fun3, color_fun4))
p3 <- plot_genes_in_pseudotime(mycds[s.genes, ], color_by = "celltype") + myth + scale_color_manual(values = c(color_fun7, color_fun3, color_fun4))
plotc <- p1 | p2 | p3
ggsave(filename = paste0(opt$od, "/Dispgenes_top6_visual_by_celltype.pdf"), plot = plotc, width = 10, height = 8)

## 6.pseudotime-dependent genes：拟时相关基因聚类热图
diff_test <- differentialGeneTest(mycds[ordering_genes, ], cores = opt$ncpu, fullModelFormulaStr = "~sm.ns(Pseudotime)")
diff_test_sig <- subset(diff_test, qval < opt$pseu_qval)
write.csv(diff_test, paste0(opt$od, "/Pseudotime_Diff_test.csv"))
write.csv(diff_test_sig, paste0(opt$od, "/Pseudotime_Diff_test_sig.csv"))

# sig_gene_names <- row.names(diff_test_sig)
sig_gene_names <- unique(diff_test_sig$gene_short_name)

# visualize modules of genes that co-vary across pseudotime
# p2 <- plot_pseudotime_heatmap(mycds[sig_gene_names, ], num_clusters = length(unique(pData(mycds)$State)), show_rownames = T, return_heatmap = T)
p2 <- plot_pseudotime_heatmap(mycds[sig_gene_names, ], num_clusters = length(unique(scRNA$celltype)), show_rownames = F, return_heatmap = T)
ggsave(filename = paste0(opt$od, "/Pseudotime_heatmap.pdf"), plot = p2, width = 8, height = 16)

# get cluster info
row_cluster <- cutree(p2$tree_row, k = length(unique(scRNA$celltype)))
write.csv(row_cluster, file = paste0(opt$od, "/Pseudotime_gene_clusters.csv"))

# 展示指定基因
if (!is.null(opt$xgene)) {
  xgene <- read.table(opt$xgene, header = F, sep = "\t", comment.char = "", check.names = F) %>% pull(1)
  # #
  # p1 <- plot_genes_jitter(mycds[xgene, ], grouping = "State", color_by = "State") + myth + scale_color_manual(values = color_fun3)
  # p2 <- plot_genes_violin(mycds[xgene, ], grouping = "State", color_by = "State") + myth + scale_fill_manual(values = color_fun3)
  # p3 <- plot_genes_in_pseudotime(mycds[xgene, ], color_by = "State") + myth + scale_color_manual(values = color_fun3)
  # plotc <- p1 | p2 | p3
  # ggsave(filename = paste0(opt$od, "/Dispgenes_xgene_visual_by_State.pdf"), plot = plotc, width = 8, height = length(xgene) * 1.5)
  # #
  # p1 <- plot_genes_jitter(mycds[xgene, ], grouping = "celltype", color_by = "celltype") + myth + scale_color_manual(values = c(color_fun7, color_fun3, color_fun4))
  # p2 <- plot_genes_violin(mycds[xgene, ], grouping = "celltype", color_by = "celltype") + myth + scale_fill_manual(values = c(color_fun7, color_fun3, color_fun4))
  # p3 <- plot_genes_in_pseudotime(mycds[xgene, ], color_by = "celltype") + myth + scale_color_manual(values = c(color_fun7, color_fun3, color_fun4))
  # plotc <- p1 | p2 | p3
  # ggsave(filename = paste0(opt$od, "/Dispgenes_xgene_visual_by_celltype.pdf"), plot = plotc, width = 10, height = length(xgene) * 1.5)
  #
  p <- plot_pseudotime_heatmap(mycds[xgene, ], num_clusters = length(unique(pData(mycds)$celltype)), show_rownames = T, return_heatmap = T)
  ggsave(filename = paste0(opt$od, "/Pseudotime_xgene_heatmap.pdf"), plot = p, width = 6, height = 12)
}


## branched expression analysis modeling: BEAM分析
# beam_res <- BEAM(mycds, branch_point = 1, cores = opt$ncpu)
# beam_res <- beam_res[order(beam_res$qval), ]
# write.csv(beam_res, file=paste0(opt$od,"/BEAM_res.csv") )
#
# table(beam_res$qval < 1e-6)
#
# if( length(which(beam_res$qval < 1e-4)) >150) {
#   mycds_sub_beam <- mycds[1:150, ]
# }
#
# options(bitmapType = "cairo")
# pdf(file = paste0(opt$od,"/BEAM_branched_heatmap.pdf"), width = 8, height = 13)
# p <- plot_genes_branched_heatmap(mycds_sub_beam,  branch_point = 1, num_clusters = length(unique(scRNA$celltype)),
#                             show_rownames = T, return_heatmap = T)
# dev.off()


# 分支太多解决方法:
# https://github.com/cole-trapnell-lab/monocle-release/issues/377
# In reduceDimension() function, there are a few arguements that can affect braching:
# 1.residualModelFormulaStr, this arguement allows you to minimize effects of some unwanted factors, e.g., from batches, treatments, etc. Sounds like regression :）
# 2.auto_param_selection, if you set it FALSE, you gotta customize some parameters passed from DDRTree() function (also developed by the monocle Team). In paremeters from this funtion, you can specify ncenter, sigma, lambda parameters, which can affect branches. But you gotta adjust them repeatedly to optimize the trace.
# I hope this would be helpful! :)
